'''
Created on Nov 13, 2013

@author: yniknafs
'''
import pymongo
import numpy as np
from ssea.countdata import BigCountMatrix

NUM_TESTS = 10000
NUM_TESTS_SR = 110
NUM_ROWS = 380000
NUM_COLS = 6300
SEED = 123987510

def db_connect(name):
    client = pymongo.MongoClient()
    db = client['eit']
    coll = db[name]
    return coll


def query_test_one():
    rand_row = np.random.randint(0, NUM_ROWS)
    rand_col = np.random.randint(0, NUM_COLS)
    coll = db_connect('bigdaddy')
    print coll.find_one({'t_id':rand_row, 's_id':rand_col}, {'_id':0})

def query_test_bigdaddy():
    coll = db_connect('bigdaddy')
    tot = 0
    np.random.seed(SEED)
    for x in xrange(NUM_TESTS):
        rand_row = np.random.randint(0, NUM_ROWS)
        rand_col = np.random.randint(0, NUM_COLS)
        v = coll.find_one({'t_id':rand_row, 's_id':rand_col}, {'value':1})['value']
        tot += v
    print tot

def query_test_mongo_index():
    coll = db_connect('fiftymill')
    tot = 0
    np.random.seed(SEED)
    for x in xrange(NUM_TESTS):
        rand_row = np.random.randint(0, NUM_ROWS)
        rand_col = np.random.randint(0, NUM_COLS)
        v = coll.find_one({'t_id':rand_row, 's_id':rand_col}, {'value':1})['value']
        tot += v
    print tot
def query_test_bigdaddy_batch():
    coll = db_connect('bigdaddy')
    tot = 0
    np.random.seed(SEED)
    rows = []
    cols = []
    for x in xrange(NUM_TESTS_SR):
        rows.append(np.random.randint(0, NUM_ROWS))
        cols.append(np.random.randint(0, NUM_COLS))
    rows.sort()
    cols.sort()
    docs = coll.find({'t_id': {'$in': rows}, 's_id': {'$in':cols}}, {'value':1})
    i = 0
    for doc in docs:
        v = doc['value']
        tot +=v
        i+=1
    print 'num docs: %d' %  i


    
    
    
    
    #rand_row = np.random.randint(0, NUM_ROWS)
    #rand_col = np.random.randint(0, NUM_COLS)
    #print coll.find_one({'t_id':rand_row, 's_id':rand_col}, {'_id':0})
    
def query_test_numpy():
    bcm = BigCountMatrix.open('/home/yniknafs/ssea/isoform_count_matrix_v4/')
    tot = 0
    np.random.seed(SEED)
    for x in xrange(NUM_TESTS):
        rand_row = np.random.randint(0, NUM_ROWS)
        rand_col = np.random.randint(0, NUM_COLS)
        v = bcm.counts[rand_row,rand_col]
        tot += v
    print 'num docs: %d' % NUM_TESTS
    # row_len, col_len = bcm.shape
    #print matrix[rand_row, rand_col]
    
if __name__ == "__main__":
    import timeit
    print 'numpy'
    print timeit.timeit('query_test_numpy()', setup='from __main__ import query_test_numpy', number=1)
    #print 'mongo_index'
    #print timeit.timeit('query_test_mongo_index()', setup='from __main__ import query_test_mongo_index', number=1)
    print 'mongo_bigdaddy'
    print timeit.timeit('query_test_bigdaddy_batch()', setup='from __main__ import query_test_bigdaddy_batch', number=1)
